import numpy as np
import matplotlib.pyplot as plt
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error, r2_score
import pandas as pd
from scipy import stats
import argparse
import logging
import os

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler("pascal_triangle.log"),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC", "Microsoft YaHei"]
plt.rcParams['axes.unicode_minus'] = False

def generate_pascal_triangle(rows, validate=True):
    if rows < 0:
        raise ValueError("杨辉三角的行数不能为负数")
    
    triangle = []
    for i in range(rows):
        row = [1] * (i + 1)
        for j in range(1, i):
            row[j] = triangle[i-1][j-1] + triangle[i-1][j]
        triangle.append(row)
    
    if validate:
        validate_pascal_triangle(triangle)
    
    return triangle

def validate_pascal_triangle(triangle):
    for i, row in enumerate(triangle):
        if len(row) != i + 1:
            raise ValueError(f"第{i}行的长度应为{i+1}，但实际为{len(row)}")
        
        if row[0] != 1 or row[-1] != 1:
            raise ValueError(f"第{i}行的首尾元素应为1，但实际为{row[0]}和{row[-1]}")
        
        for j in range(1, i):
            if row[j] != triangle[i-1][j-1] + triangle[i-1][j]:
                raise ValueError(f"第{i}行第{j}个元素不符合递推关系")
    
    logger.info(f"杨辉三角验证通过，共{len(triangle)}行")

def create_dataset(triangle, add_polynomial_features=False):
    X = []
    y = []
    
    for i, row in enumerate(triangle):
        for j, value in enumerate(row):
            center_distance = abs(j - i/2)
            features = [i, j, center_distance]
            
            if add_polynomial_features:
                features.extend([i**2, j**2, center_distance**2, i*j])
            
            X.append(features)
            y.append(value)
    
    return np.array(X), np.array(y)

def visualize_pascal_triangle(triangle, title="杨辉三角可视化", save_path="pascal_triangle_visualization.png"):
    if not triangle:
        logger.warning("杨辉三角为空，无法进行可视化")
        return
    
    rows = len(triangle)
    max_width = len(triangle[-1]) * 2 - 1
    
    fig_width = min(30, max(10, rows // 1.5))
    fig_height = min(25, max(8, rows // 1.8))
    
    fig, ax = plt.subplots(figsize=(fig_width, fig_height))
    
    if rows <= 15:
        font_size = min(14, 180 // rows)
        point_size = min(300, 4000 // rows)
    elif rows <= 30:
        font_size = min(10, 150 // rows)
        point_size = min(200, 3000 // rows)
    elif rows <= 50:
        font_size = min(7, 120 // rows)
        point_size = min(120, 2000 // rows)
    else:
        font_size = min(5, 100 // rows)
        point_size = min(80, 1500 // rows)
    
    for i, row in enumerate(triangle):
        x_positions = np.linspace(0, max_width, len(row))
        log_vals = [np.log(v + 1) for v in row]
        scatter = ax.scatter(x_positions, [-i]*len(row), s=point_size, 
                             c=log_vals, cmap='viridis', alpha=0.9, edgecolors='w', linewidths=0.5)
        
        for j, val in enumerate(row):
            ax.annotate(str(val), (x_positions[j], -i), ha='center', va='center', 
                        color='white', fontweight='bold', fontsize=font_size)
    
    cbar = plt.colorbar(scatter)
    cbar.set_label('数值对数缩放')
    
    ax.set_title(title, fontsize=16)
    ax.set_axis_off()
    plt.tight_layout()
    
    try:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        logger.info(f"杨辉三角可视化已保存至 {save_path}")
    except Exception as e:
        logger.error(f"保存杨辉三角可视化失败: {e}")
    
    plt.show()

def visualize_pascal_triangle_blocks(triangle, block_size=15, title="杨辉三角分块可视化", save_dir="pascal_blocks"):
    rows = len(triangle)
    if rows <= block_size:
        logger.info("杨辉三角行数少于分块大小，使用标准可视化模式")
        visualize_pascal_triangle(triangle, title, f"{save_dir}.png")
        return
    
    os.makedirs(save_dir, exist_ok=True)
    
    num_blocks = (rows + block_size - 1) // block_size
    
    logger.info(f"将杨辉三角分为{num_blocks}个块进行可视化")
    
    for block_idx in range(num_blocks):
        start_row = block_idx * block_size
        end_row = min((block_idx + 1) * block_size, rows)
        
        block_triangle = triangle[start_row:end_row]
        
        fig_width = min(30, max(10, len(block_triangle) // 1.5))
        fig_height = min(25, max(8, len(block_triangle) // 1.8))
        
        fig, ax = plt.subplots(figsize=(fig_width, fig_height))
        
        if len(block_triangle) <= 15:
            font_size = min(14, 180 // len(block_triangle))
            point_size = min(300, 4000 // len(block_triangle))
        elif len(block_triangle) <= 30:
            font_size = min(10, 150 // len(block_triangle))
            point_size = min(200, 3000 // len(block_triangle))
        else:
            font_size = min(7, 120 // len(block_triangle))
            point_size = min(120, 2000 // len(block_triangle))
        
        max_width = len(block_triangle[-1]) * 2 - 1 if block_triangle else 0
        
        for i, row in enumerate(block_triangle):
            rel_row = i + start_row
            x_positions = np.linspace(0, max_width, len(row))
            log_vals = [np.log(v + 1) for v in row]
            scatter = ax.scatter(x_positions, [-i]*len(row), s=point_size, 
                                 c=log_vals, cmap='viridis', alpha=0.9, edgecolors='w', linewidths=0.5)
            
            for j, val in enumerate(row):
                ax.annotate(str(val), (x_positions[j], -i), ha='center', va='center', 
                            color='white', fontweight='bold', fontsize=font_size)
        
        block_title = f"{title} (行 {start_row}-{end_row-1})"
        ax.set_title(block_title, fontsize=16)
        ax.set_axis_off()
        
        if block_idx > 0:
            ax.text(-5, 0, f"↑ 上接行 {start_row-1}", fontsize=8, color='gray')
        if block_idx < num_blocks - 1:
            ax.text(-5, -len(block_triangle)+1, f"↓ 下接行 {end_row}", fontsize=8, color='gray')
        
        plt.tight_layout()
        
        save_path = os.path.join(save_dir, f"block_{block_idx}.png")
        try:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            logger.info(f"杨辉三角块 {block_idx} 已保存至 {save_path}")
        except Exception as e:
            logger.error(f"保存杨辉三角块 {block_idx} 失败: {e}")
        
        plt.close()
    
    logger.info(f"所有杨辉三角块已保存至目录: {save_dir}")

def visualize_pascal_triangle_medium(triangle, title="杨辉三角可视化（中型）", save_path="pascal_triangle_medium.png"):
    if not triangle:
        logger.warning("杨辉三角为空，无法进行可视化")
        return
    
    rows = len(triangle)
    if rows < 20 or rows > 50:
        logger.warning(f"当前行数为{rows}，推荐使用标准或分块可视化模式")
        visualize_pascal_triangle(triangle, title, save_path)
        return
    
    max_width = len(triangle[-1]) * 2 - 1
    
    fig_width = min(40, max(15, rows))
    fig_height = min(30, max(10, rows * 0.7))
    
    fig, ax = plt.subplots(figsize=(fig_width, fig_height))
    
    font_size = max(4, min(9, 200 // rows))
    point_size = max(60, min(180, 4000 // rows))
    
    for i, row in enumerate(triangle):
        x_positions = np.linspace(0, max_width, len(row))
        log_vals = [np.log(v + 1) for v in row]
        scatter = ax.scatter(x_positions, [-i]*len(row), s=point_size, 
                             c=log_vals, cmap='viridis', alpha=0.9, edgecolors='w', linewidths=0.5)
        
        for j, val in enumerate(row):
            ax.annotate(str(val), (x_positions[j], -i), ha='center', va='center', 
                        color='white', fontweight='bold', fontsize=font_size)
    
    cbar = plt.colorbar(scatter)
    cbar.set_label('数值对数缩放')
    
    ax.set_title(title, fontsize=18)
    ax.set_axis_off()
    plt.tight_layout()
    
    try:
        plt.savefig(save_path, dpi=400, bbox_inches='tight')
        logger.info(f"中型杨辉三角可视化已保存至 {save_path}")
    except Exception as e:
        logger.error(f"保存中型杨辉三角可视化失败: {e}")
    
    plt.show()

def build_prediction_model(X, y, test_size=0.2, random_state=42):
    if X.shape[0] == 0:
        raise ValueError("数据集为空，无法构建模型")
    
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=test_size, random_state=random_state)
    
    model = LinearRegression()
    model.fit(X_train, y_train)
    
    y_pred = model.predict(X_test)
    mse = mean_squared_error(y_test, y_pred)
    r2 = r2_score(y_test, y_pred)
    
    logger.info(f"模型评估：MSE={mse:.4f}, R²={r2:.4f}")
    
    feature_names = ['行号', '列号', '到中心距离']
    if X.shape[1] > 3:
        feature_names.extend(['行号²', '列号²', '到中心距离²', '行号×列号'])
    
    feature_importance = pd.DataFrame({
        '特征': feature_names,
        '重要性': model.coef_
    }).sort_values('重要性', ascending=False)
    
    return model, y_pred, {'mse': mse, 'r2': r2, 'feature_importance': feature_importance}

def discover_patterns(triangle):
    if not triangle:
        logger.warning("杨辉三角为空，无法发现数学规律")
        return
    
    logger.info("\n=== 杨辉三角中的数学规律 ===")
    
    rows_to_check = [min(10, len(triangle)-1)]
    logger.info("\n1. 二项式系数验证：")
    for n in rows_to_check:
        binomial_coeffs = [int(np.math.comb(n, k)) for k in range(n+1)]
        logger.info(f"第{n}行（二项式系数）: {binomial_coeffs}")
        logger.info(f"杨辉三角对应行: {triangle[n]}")
        logger.info(f"验证结果: {binomial_coeffs == triangle[n]}\n")
    
    logger.info("\n2. 斐波那契数列验证：")
    max_diagonals = min(10, len(triangle))
    fibonacci = []
    for k in range(max_diagonals):
        diag_sum = sum(triangle[i][j] for i in range(len(triangle)) 
                     for j in range(len(triangle[i])) 
                     if i - j == k)
        fibonacci.append(diag_sum)
    
    logger.info(f"从对角线计算的斐波那契数列: {fibonacci}")
    
    logger.info("\n3. 奇偶分布规律：")
    parity_counts = []
    for i, row in enumerate(triangle):
        odds = sum(1 for num in row if num % 2 != 0)
        parity_counts.append((i, odds, len(row) - odds))
    
    logger.info("前几行的奇偶分布（行号, 奇数个数, 偶数个数）:")
    for i, o, e in parity_counts[:min(10, len(parity_counts))]:
        logger.info(f"行 {i}: 奇数={o}, 偶数={e}")
    
    visualize_parity_pattern(triangle)

def visualize_parity_pattern(triangle, save_path="sierpinski_triangle.png"):
    if not triangle:
        logger.warning("杨辉三角为空，无法可视化奇偶分布")
        return
    
    max_rows = len(triangle)
    data = np.zeros((max_rows, max_rows*2 - 1))
    
    for i, row in enumerate(triangle):
        start_idx = (max_rows*2 - 1 - len(row)) // 2
        for j, num in enumerate(row):
            data[i, start_idx + j] = num % 2
    
    plt.figure(figsize=(12, 8))
    plt.imshow(data, cmap='binary', interpolation='nearest')
    plt.title('杨辉三角奇偶分布（Sierpinski三角形）')
    plt.axis('off')
    
    try:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        logger.info(f"奇偶分布可视化已保存至 {save_path}")
    except Exception as e:
        logger.error(f"保存奇偶分布可视化失败: {e}")
    
    plt.show()

def predict_new_values(model, n, k, feature_names=None):
    if n < 0 or k < 0 or k > n:
        raise ValueError(f"无效的位置: 行号={n}, 列号={k}。必须满足 0 <= k <= n")
    
    center_distance = abs(k - n/2)
    
    if feature_names and len(feature_names) > 3:
        features = [n, k, center_distance, n**2, k**2, center_distance**2, n*k]
    else:
        features = [n, k, center_distance]
    
    prediction = model.predict([features])[0]
    actual = int(np.math.comb(n, k))
    
    result = {
        'n': n,
        'k': k,
        'prediction': prediction,
        'actual': actual,
        'relative_error': abs(prediction - actual)/actual * 100 if actual != 0 else 0
    }
    
    logger.info(f"预测杨辉三角中第{n}行第{k}列的值: 预测={prediction:.4f}, 实际={actual}, 相对误差={result['relative_error']:.2f}%")
    return result

def check_chinese_fonts():
    available_fonts = plt.rcParams["font.family"]
    logger.info(f"当前设置的中文字体: {available_fonts}")
    
    try:
        plt.figure(figsize=(4, 2))
        plt.title("中文字体测试", fontsize=12)
        plt.axis('off')
        plt.close()
        logger.info("中文字体测试成功")
        return True
    except Exception as e:
        logger.warning(f"中文字体测试失败: {e}")
        logger.warning("图表中的文字可能显示为方块")
        return False

def main(args):
    try:
        check_chinese_fonts()
        
        logger.info(f"正在生成杨辉三角，行数={args.rows}...")
        triangle = generate_pascal_triangle(args.rows)
        
        logger.info("\n正在可视化杨辉三角...")
        if args.rows <= 20:
            visualize_pascal_triangle(triangle)
        elif args.rows <= 50:
            visualize_pascal_triangle_medium(triangle)
        else:
            visualize_pascal_triangle_blocks(triangle, block_size=args.block_size)
        
        logger.info("\n正在准备数据集...")
        X, y = create_dataset(triangle, args.polynomial_features)
        
        logger.info("\n正在构建预测模型...")
        model, _, metrics = build_prediction_model(X, y)
        
        logger.info("\n特征重要性分析：")
        logger.info(metrics['feature_importance'])
        
        logger.info("\n正在发现数学规律...")
        discover_patterns(triangle)
        
        test_n, test_k = args.test_row, args.test_col
        if test_n >= args.rows or test_k > test_n:
            logger.info(f"\n预测超出训练范围的数据：第{test_n}行第{test_k}列")
        else:
            logger.info(f"\n预测训练范围内的数据验证：第{test_n}行第{test_k}列")
        
        feature_names = ['行号', '列号', '到中心距离']
        if args.polynomial_features:
            feature_names.extend(['行号²', '列号²', '到中心距离²', '行号×列号'])
        
        predict_new_values(model, test_n, test_k, feature_names)
        
        logger.info("\n程序执行完毕")
        
    except Exception as e:
        logger.error(f"程序执行失败: {e}", exc_info=True)
        raise

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='杨辉三角数据挖掘分析')
    parser.add_argument('--rows', type=int, default=30, help='杨辉三角的行数')
    parser.add_argument('--polynomial_features', action='store_true', help='是否添加多项式特征')
    parser.add_argument('--test_row', type=int, default=35, help='测试预测的行号')
    parser.add_argument('--test_col', type=int, default=17, help='测试预测的列号')
    parser.add_argument('--block_size', type=int, default=20, help='分块可视化时每块的行数')
    
    args = parser.parse_args()
    
    if args.test_col > args.test_row:
        parser.error(f"测试列号({args.test_col})不能大于测试行号({args.test_row})")
    
    main(args)    